//
//  main.m
//  SlimMachO
//
//  Created by  JasonLee on 2022/1/10.
//

#import <Foundation/Foundation.h>
#include <mach-o/loader.h>
#include <mach-o/fat.h>
#include <sys/stat.h>

#define IS_64_BIT(x) ((x) == MH_MAGIC_64 || (x) == MH_CIGAM_64)
#define IS_LITTLE_ENDIAN(x) ((x) == FAT_CIGAM || (x) == MH_CIGAM_64 || (x) == MH_CIGAM)
#define SWAP32(x, magic) (IS_LITTLE_ENDIAN(magic)? OSSwapInt32(x): (x))
#define SWAP64(x, magic) (IS_LITTLE_ENDIAN(magic)? OSSwapInt64(x): (x))

#define ROUND_UP(x, y) (((x) + (y) - 1) & -(y))

#define ABSDIFF(x, y) ((x) > (y)? (uintmax_t)(x) - (uintmax_t)(y): (uintmax_t)(y) - (uintmax_t)(x))

#define BUFSIZE 512

#define LC_COMMAND_NAME "LC_LOAD_DYLIB"

void fbzero(FILE *f, off_t offset, size_t len) {
    static unsigned char zeros[BUFSIZE] = {0};
    fseeko(f, offset, SEEK_SET);
    while(len != 0) {
        size_t size = MIN(len, sizeof(zeros));
        fwrite(zeros, size, 1, f);
        len -= size;
    }
}
size_t fpeek(void *ptr, size_t size, size_t nitems, FILE *stream) {
    off_t pos = ftello(stream);
    size_t result = fread(ptr, size, nitems, stream);
    fseeko(stream, pos, SEEK_SET);
    return result;
}
void *read_load_command(FILE *f, uint32_t cmdsize) {
    void *lc = malloc(cmdsize);

    fpeek(lc, cmdsize, 1, f);

    return lc;
}

const char hex_asc[] = "0123456789abcdef";
#define hex_asc_lo(x)    hex_asc[((x) & 0x0f)]
#define hex_asc_hi(x)    hex_asc[((x) & 0xf0) >> 4]
void hex_dump_to_buffer(const void *buf, size_t len, char *linebuf, size_t linebuflen)
{
    const char *ptr = (const char *)buf;
    int lx = 0;
    int ascii_column;

    if (!len) {
        goto end;
    } else if (len > 16) {
        len = 16;
    }
    int j;
    for (j = 0; (j < len) && (lx + 3) <= linebuflen; j++) {
        char ch = ptr[j];
        linebuf[lx++] = hex_asc_hi(ch);
        linebuf[lx++] = hex_asc_lo(ch);
        linebuf[lx++] = ' ';
    }
    if (j) {
        lx--;
    }
    ascii_column = 3 * 16 + 2;
    
    while (lx < (linebuflen - 1) && lx < (ascii_column - 1))
        linebuf[lx++] = ' ';
    for (j = 0; (j < len) && (lx + 2) < linebuflen; j++) {
        char ch = ptr[j];
        linebuf[lx++] = (isascii(ch) && isprint(ch)) ? ch : '.';
    }
    
end:
    linebuf[lx++] = '\0';
}

void print_hex_dump_bytes(const void *buf, size_t len)
{
    const char *ptr = (const char *)buf;
    unsigned long linelen;
    size_t remaining = len;
    
    char linebuf[16 * 3 + 2 + 32 + 1];
    
    for (int i = 0; i < len; i += 16) {
        linelen = MIN(remaining, 16);
        remaining -= 16;
        hex_dump_to_buffer(ptr + i, linelen, linebuf, sizeof(linebuf));
        NSLog(@"%s\n", linebuf);
    }
}
BOOL isMachO(const char * file ){
    FILE *f = fopen(file, "r");
    NSLog(@"%d",errno);
    if (!f) {
        return false;
    }
    rewind(f);
    uint32_t magic;
    fread(&magic, sizeof(uint32_t), 1, f);
    return magic == FAT_MAGIC || magic == FAT_CIGAM || magic == MH_MAGIC || magic == MH_CIGAM || magic == MH_MAGIC_64 || magic == MH_CIGAM_64;
}
int slim_arm64e_binary(const char *file)
{
    int ret = 0;
    FILE *f = fopen(file, "r");
    if (!f) {
        printf("fail to open %s, %d", file, errno);
        return -1;
    }

    uint32_t magic;
    fread(&magic, sizeof(uint32_t), 1, f);
//    if ((magic != FAT_MAGIC) &&
//        (magic != FAT_CIGAM)) {
//        return ret;
//    }

    rewind(f);

    struct fat_header fh;
    fread(&fh, sizeof(fh), 1, f);
    uint32_t nfat_arch = SWAP32(fh.nfat_arch, magic);

    struct fat_arch archs[nfat_arch];
    fread(archs, sizeof(archs), 1, f);
    
    char tmp[1024] = {0};
    sprintf(tmp, "%s.tmp", file);
    FILE *tmpf = fopen(tmp, "w+");
    struct stat state;
    stat(file, &state);
    //tmpf 设置为 跟f一样大小
    ftruncate(fileno(tmpf), state.st_size);
    rewind(tmpf);
    //将fatheader写入tmp
    fwrite(&fh, sizeof(fh), 1, tmpf);
    //获取x86原来offset
    uint32_t offset = SWAP32(archs[0].offset, magic);
    //获取x86原来总大小
    uint32_t size = SWAP32(archs[0].size, magic);
    
    ftruncate(fileno(tmpf), size);

    //修改 offset
    archs[0].offset = SWAP32(nfat_arch * sizeof(struct fat_arch ) + sizeof(fh), magic);
    archs[0].align = SWAP32(0, magic);
    
    uint32_t distance = offset - nfat_arch * sizeof(struct fat_arch ) -+ sizeof(fh);
    //将fat_arch写入tmpf
    fwrite(&archs, nfat_arch * sizeof(struct fat_arch) , 1, tmpf);
  //读取mach_header
    fseeko(f, offset, SEEK_SET);
    struct mach_header header;
    fread(&header, sizeof(header), 1, f);
    
//
//    fwrite(&header, sizeof(header) , 1, tmpf);
    
    uint32_t cmdSize =  SWAP32(header.sizeofcmds, header.magic);
    uint32_t ncmds =  SWAP32(header.ncmds, header.magic);
    
    //读取mach_header loadcommand
    char *buf = (char *)malloc(cmdSize + sizeof(struct mach_header));
    if (!buf) {
        return -1;
    }
    fseeko(f, offset, SEEK_SET);
    fread(buf, cmdSize + sizeof(struct mach_header) , 1, f);
    //将mach_header + loadCommand 写入tmpf
    fwrite(buf, cmdSize + sizeof(struct mach_header), 1, tmpf);
    
    fseeko(f, offset + (IS_64_BIT(header.magic)? sizeof(struct mach_header_64): sizeof(struct mach_header)),SEEK_SET);
    fseeko(tmpf, offset + (IS_64_BIT(header.magic)? sizeof(struct mach_header_64): sizeof(struct mach_header)),SEEK_SET);

    off_t symtab_pos = -1;
    uint32_t symtab_size = 0;
    for (int i = 0; i < ncmds; i++) {
        struct load_command lc;
        fpeek(&lc, sizeof(lc), 1, f);
        uint32_t cmdsize = SWAP32(lc.cmdsize, header.magic);
        uint32_t cmd = SWAP32(lc.cmd, header.magic);
        switch (cmd) {
            case LC_SEGMENT:
            case LC_SEGMENT_64:
                if (cmd == LC_SEGMENT) {
                    struct segment_command *cmd = (struct segment_command *)read_load_command(f, cmdsize);
                    for (int i = 0; i < cmd->nsects; i++) {
                        struct section sectionH;
                        fpeek(&sectionH, sizeof(sectionH), 1, f);
                        if (sectionH.size != 0) {
                            sectionH.offset = sectionH.offset + distance;
                        }
                        fwrite(&sectionH, sizeof(sectionH), 1, tmpf);
                    }

                    free(cmd);
                } else {
                    struct segment_command_64 *cmd = (struct segment_command_64 *)read_load_command(f, cmdsize);
                    off_t pos = ftello(f);
                    fseek(f, pos + sizeof(struct segment_command_64), SEEK_SET);
                    if (!strcmp("__LINKEDIT", cmd->segname)) {
                        cmd->fileoff = cmd->fileoff + distance;
                        fwrite(cmd, sizeof(struct segment_command_64), 1, tmpf);
                    } else {
                        for (int i = 0; i < SWAP32(cmd->nsects, header.magic); i++) {
                            struct section_64 sectionH;
                            fpeek(&sectionH, sizeof(sectionH), 1, f);
                            off_t curpos = ftello(f);
                            fseeko(f, sizeof(sectionH), SEEK_CUR);
                            
                            if (sectionH.size != 0) {
                                sectionH.offset = sectionH.offset + distance;
                            }
                            fseeko(tmpf, curpos - distance, SEEK_SET);
                            fwrite(&sectionH, sizeof(sectionH), 1, tmpf);
                        }
                    }
                    
                    fseek(f, pos, SEEK_SET);
                    free(cmd);
                }
                break;
            case  LC_DYLD_INFO_ONLY: {
                struct dyld_info_command *dyldinfo = (struct dyld_info_command *)read_load_command(f, cmdsize);
                dyldinfo->rebase_off = dyldinfo->rebase_off + distance;
                dyldinfo->bind_off = dyldinfo->bind_off + distance;
                dyldinfo->weak_bind_off = dyldinfo->weak_bind_off + distance;
                dyldinfo->lazy_bind_off = dyldinfo->lazy_bind_off + distance;
                dyldinfo->export_off = dyldinfo->export_off + distance;
                off_t curpos = ftello(f);
                
                fseeko(tmpf, curpos - distance, SEEK_SET);
                fwrite(dyldinfo, sizeof(struct dyld_info_command), 1, tmpf);
                break;
            }
                
            case LC_SYMTAB: {
                struct symtab_command *symtab = (struct symtab_command *)read_load_command(f, cmdsize);
                symtab->symoff = symtab->symoff + distance;
                symtab->stroff = symtab->stroff + distance;
                off_t curpos = ftello(f);
                fseeko(tmpf, curpos - distance, SEEK_SET);
                fwrite(symtab, sizeof(struct symtab_command), 1, tmpf);
                symtab_pos = ftello(f);
                symtab_size = cmdsize;
                break;
            }
            case LC_DYSYMTAB: {
                struct dysymtab_command *dysym  = (struct dysymtab_command *)read_load_command(f, cmdsize);
                dysym->tocoff = dysym->tocoff + distance;
                dysym->modtaboff = dysym->modtaboff + distance;
                dysym->extrefsymoff = dysym->extrefsymoff + distance;
                dysym->indirectsymoff = dysym->indirectsymoff + distance;
                dysym->extreloff = dysym->extreloff + distance;
                dysym->locreloff = dysym->locreloff + distance;
                off_t curpos = ftello(f);
                fseeko(tmpf, curpos - distance, SEEK_SET);
                fwrite(dysym, sizeof(struct dysymtab_command), 1, tmpf);
                break;

            }
            case LC_MAIN: {
                struct entry_point_command *entry  = (struct entry_point_command *)read_load_command(f, cmdsize);
                entry->entryoff = entry->entryoff + distance;
                off_t curpos = ftello(f);
                fseeko(tmpf, curpos - distance, SEEK_SET);
                fwrite(entry, sizeof(struct entry_point_command), 1, tmpf);

                break;
            }
            case LC_CODE_SIGNATURE: {
                struct linkedit_data_command *linkedit = (struct linkedit_data_command *)read_load_command(f, cmdsize);
                linkedit->dataoff = linkedit->dataoff + distance;
                off_t curpos = ftello(f);
                fseeko(tmpf, curpos - distance, SEEK_SET);
                fwrite(linkedit, sizeof(struct linkedit_data_command), 1, tmpf);
                break;
            }
        }
        fseeko(f, SWAP32(lc.cmdsize, header.magic), SEEK_CUR);
    }
    
    
    long offset112 = ftell(f);
    char *sectionBuf = (char *)malloc(size - cmdSize - sizeof(struct mach_header) - 4);
    if (!sectionBuf) {
        return -1;
    }
    fread(sectionBuf, size - cmdSize - sizeof(struct mach_header), 1, f);
    fseeko(tmpf, offset112, SEEK_SET);
    fwrite(sectionBuf,  size, 1, tmpf);
//    print_hex_dump_bytes(sectionBuf, size - cmdSize - sizeof(struct mach_header) - 4);

    fclose(tmpf);
    return ret;
}
int main(int argc, const char * argv[]) {
    @autoreleasepool {
        BOOL ret = isMachO("/usr/bin/perlbug");
        NSLog(@"======ret:%d",ret);
        slim_arm64e_binary("/Users/jasonlee/Desktop/Google Chrome2");
        // insert code here...
        NSLog(@"Hello, World!");
    }
    return 0;
}









